from docplex.mp.model import Model
from network import QKDNetwork
import numpy as np
from generate import GenerateNetwork

class RMP():
    def __init__(self, length, net: QKDNetwork):
        self.length = length
        self.net = net
        self.model = Model("RMP")
        self.alpha = self.model.continuous_var(lb=0, name="Alpha")  # 1️⃣
        self.dual_lambda = {}  # 对偶约束 lambda 2️⃣
        self.dual_mu = {}  # 对偶约束 mu 3️⃣
        self.x = {}  # x_{ij}^{mn} 4️⃣
        self.p = {}  # p_{ij}^{mnuv} 5️⃣
        self.added_p = []  # 已添加到模型的P
        self.sd = []  # 业务量
        self.dual_index = 0
        self.last_value = 0
        for i in range(length):
            self.dual_lambda[i] = {}
            self.dual_mu[i] = {}
            for j in range(length):
                self.dual_lambda[i][j] = []
                self.dual_mu[i][j] = []
                for k in range(length):
                    self.dual_lambda[i][j].append(self.model.continuous_var(
                        lb=0, name=f'lambda_{i}_{j}_{k}'))
                    self.dual_mu[i][j].append(self.model.continuous_var(
                        lb=0, name=f'mu_{i}_{j}_{k}'))
        for i in range(1, length):
            for j in range(1, length):
                if i != j:
                    self.sd.append((i, j))
        for i in range(length):
            self.x[i] = {}
            for j in range(length):
                self.x[i][j] = {}
                for k in range(length):
                    self.x[i][j][k] = {}
                    for l in range(length):
                        self.x[i][j][k][l] = self.model.continuous_var(
                            lb=0, ub=1, name=f'x_{i}_{j}_{k}_{l}')
        for i in range(length):
            self.p[i] = {}
            for j in range(length):
                self.p[i][j] = {}
                for k in range(length):
                    self.p[i][j][k] = {}
                    for l in range(length):
                        self.p[i][j][k][l] = {}
                        for m in range(length):
                            self.p[i][j][k][l][m] = []
                            for q in range(length):
                                self.p[i][j][k][l][m].append(self.model.continuous_var(
                                    lb=0, ub=1, name=f'p_{i}_{j}_{k}_{l}_{m}_{q}'))

    def runCG(self):
        self.initShortestPath()
        self.buildTotalConstraint()
        self.confirmAddSingleColumn()

    def addPrimaryConstraint(self):
        self.dual_index = 0
        # 增加基础的约束（非列生成部分）
        # 约束1️⃣(1)
        for i in self.net.nodeList:
            nodeID = i.nodeID                           # --> i
            for adjacentNodeID in i.adjacentNodes:      # --> j
                dualFlowFraction = self.model.linear_expr()
                dualFlowFraction += self.net.getLinkCapacity(
                    nodeID, adjacentNodeID) * self.alpha
                for sd_p in self.sd:
                    sd_start, sd_end = sd_p             # --> m, n
                    for t in self.net.getNode(sd_start).adjacentNodes:
                        dualFlowFraction -= self.dual_lambda[nodeID][adjacentNodeID][sd_start] * \
                            self.net.getLinkCapacity(sd_start, t)
                    for s in self.net.getNode(sd_end).adjacentNodes:
                        dualFlowFraction -= self.dual_mu[nodeID][adjacentNodeID][sd_end] * \
                            self.net.getLinkCapacity(s, sd_end)
                self.model.add_constraint(dualFlowFraction >= 0)
                self.dual_index += 1

        # 约束2️⃣
        for sd_p in self.sd:
            sd_start, sd_end = sd_p  # --> m, n
            for i in self.net.nodeList:
                nodeID = i.nodeID    # --> i
                outFlow = self.model.linear_expr()
                if nodeID == sd_start:
                    outFlow -= 1
                if nodeID == sd_end:
                    outFlow += 1
                for adjacentNodeID in i.adjacentNodes:
                    outFlow += self.x[sd_start][sd_end][nodeID][adjacentNodeID]
                    outFlow -= self.x[sd_start][sd_end][adjacentNodeID][nodeID]
                self.model.add_constraint(outFlow == 0)
                self.dual_index += 1

    def buildTotalConstraint(self):
        # 1. 先清除所有的约束
        # 2. 添加初始的部分约束
        # 3. 添加列生成部分的约束
        self.model.clear_constraints()
        self.addPrimaryConstraint()  # 增加初始约束
        # 约束1️⃣(2)
        for sd_p in self.sd:
            m, n = sd_p                             # --> m, n
            for i_item in self.net.nodeList:
                i = i_item.nodeID                   # --> i
                for j in i_item.adjacentNodes:      # --> j
                    dualFraction = self.model.linear_expr()
                    dualFraction += self.dual_lambda[i][j][m]
                    dualFraction += self.dual_mu[i][j][n]
                    for v_item in self.net.nodeList:
                        v = v_item.nodeID
                        for t in v_item.adjacentNodes:
                            # dualFraction
                            if (self.existInP(m, n, v, t, i, j)):
                                dualFraction -= self.p[m][n][v][t][i][j]
                    self.model.add_constraint(dualFraction >= 0)
        for sd_p in self.sd:
            m, n = sd_p                             # m, n
            for i_item in self.net.nodeList:
                i = i_item.nodeID                   # i
                for v_item in self.net.nodeList:
                    v = v_item.nodeID               # v
                    for t in v_item.adjacentNodes:  # t
                        fractionFlow = self.model.linear_expr()
                        if (i == v):
                            fractionFlow -= self.x[m][n][v][t]
                        if (i == t):
                            fractionFlow += self.x[m][n][v][t]
                        for delta in self.net.getNode(i).adjacentNodes:
                            if (self.existInP(m, n, v, t, i, delta)):
                                fractionFlow += self.p[m][n][v][t][i][delta]
                            if (self.existInP(m, n, v, t, delta, i)):
                                fractionFlow -= self.p[m][n][v][t][delta][i]
                        self.model.add_constraint(fractionFlow == 0)
        self.model.minimize(self.alpha)

    # def getObliviousRoutingScheme(self):
    #     if self.sol:
    #         x_ret = np.zeros(
    #             (self.length, self.length, self.length, self.length))
    #         p_ret = np.zeros((self.length, self.length, self.length,
    #                          self.length, self.length, self.length))
    #         for i in range(self.length):
    #             for j in range(self.length):
    #                 for k in range(self.length):
    #                     for l in range(self.length):
    #                         if self.sol[self.x[i][j][k][l]]:
    #                             x_ret[i][j][k][l] = 1
    #         for i in range(self.length):
    #             for j in range(self.length):
    #                 for k in range(self.length):
    #                     for l in range(self.length):
    #                         for m in range(self.length):
    #                             for q in range(self.length):
    #                                 value = self.sol[self.p[i][j][k][l][m][q]]
    #                                 if value != 0:
    #                                     p_ret[i][j][k][l][m][q] = value
    #         return x_ret, p_ret

    def getDualVariables(self):
        # 获取对偶变量的值
        # \varphi_{mnij}
        # \psi_{mnivt}
        dual_index_tmp = self.dual_index
        mn_length = len(self.sd) + 1
        ij_length = len(self.net.nodeList) + 1
        vt_length = len(self.net.nodeList) + 1
        self.sol = self.model.solve()
        if self.sol:
            obj = self.sol.get_objective_value()
            # print(obj)
            if (self.last_value == 0):
                self.last_value = obj
                print("目标值：", obj)
            elif (self.last_value != obj):
                print("目标值：", obj)
                self.last_value = obj

            cpx = self.model.get_cplex()
            dual = cpx.solution.get_dual_values()
            varphi_mnij = np.zeros(
                (mn_length, mn_length, ij_length, ij_length))
            psi_mnivt = np.zeros(
                (mn_length, mn_length, ij_length, vt_length, vt_length))
            for sd_p in self.sd:
                m, n = sd_p                             # m, n
                for i_item in self.net.nodeList:
                    i = i_item.nodeID                   # i
                    for j in i_item.adjacentNodes:      # j
                        varphi_mnij[m][n][i][j] = dual[dual_index_tmp]
                        dual_index_tmp += 1

            for sd_p in self.sd:
                m, n = sd_p                             # m, n
                for i_item in self.net.nodeList:
                    i = i_item.nodeID                   # i
                    for v_item in self.net.nodeList:
                        v = v_item.nodeID               # v
                        for t in v_item.adjacentNodes:  # t
                            psi_mnivt[m][n][i][v][t] = dual[dual_index_tmp]
                            dual_index_tmp += 1

            file_path = "dual.txt"
            with open(file_path, "w") as file:
                for idx in range(len(dual)):
                    line = dual[idx]
                    file.write(str(self.model.get_constraint_by_index(
                        idx)) + "：" + str(line) + "\n")

            return varphi_mnij, psi_mnivt

    def confirmAddSingleColumn(self):
        for sd_p in self.sd:
            m, n = sd_p
            for i_item in self.net.nodeList:
                i = i_item.nodeID
                for j in i_item.adjacentNodes:
                    for v_item in self.net.nodeList:
                        v = v_item.nodeID
                        for t in v_item.adjacentNodes:
                            if (not self.existInP(m, n, i, j, v, t)):
                                varphi, psi = self.getDualVariables()
                                reduced_cost = varphi[m][n][i][j] - \
                                    psi[m][n][i][v][t] + psi[m][n][j][v][t]
                                if (reduced_cost < 0):
                                    print(reduced_cost)
                                    self.addSingleColumnConstraint(
                                        [m, n, i, j, v, t])
                                    self.buildTotalConstraint()
                                else:
                                    print(f"Pass：{reduced_cost}", [m, n, i, j, v, t])

    def initShortestPath(self):
        self.added_p = []
        # 初始时计算最短路
        for sd_p in self.sd:
            m, n = sd_p
        #     for v_item in self.net.nodeList:
        #         v = v_item.nodeID
        #         for t in v_item.adjacentNodes:
        #             for i_item in self.net.nodeList:
        #                 i = i_item.nodeID
        #                 for j in i_item.adjacentNodes:
        #                     self.added_p.append([m, n, v, t, i, j])
        # self.added_p = self.added_p[50:]
            tmpShortestPath = self.net.getShortedPath(m, n)
            tmpShortestPath.insert(0, m)
            for k in range(len(tmpShortestPath) - 1):
                # 使用Shortest算法计算的最短路径，并将p[m][n][i][j][i][j]设为1
                i = tmpShortestPath[k]
                j = tmpShortestPath[k + 1]
                self.added_p.append([m, n, i, j, i, j])

    def existInP(self, a, b, c, d, e, f) -> bool:
        for p_tmp in self.added_p:
            if (a == p_tmp[0] and b == p_tmp[1] and c == p_tmp[2] and d == p_tmp[3] and e == p_tmp[4] and f == p_tmp[5]):
                return True
        return False

    def addSingleColumnConstraint(self, new_p: list):
        # new_p [1,2,3,4,5,6] 新增的p
        # 添加一列p_{ij}^{mnuv}
        print("Add Column")
        print(new_p)
        self.added_p.append(new_p)  # 将新p添加到旧有的p中


GenerateNetwork.generate_network(5, 1.5, .1, 1)
qkd_network = QKDNetwork("generate.txt")
rmp = RMP(len(qkd_network.nodeList) + 1, qkd_network)
rmp.runCG()
